"""
@File    :   orm_init_drop.py
@Time    :   2024/02/11 21:47:46
@Author  :   Qrj
@Email   :   923207736@qq.com
@description   :   orm初始化等操作
"""

import shutil

from sqlalchemy import create_engine

from alembic import command
from alembic.config import Config
from src.config import get_app_settings
from src.utils.auto_import import AutoImportUtil
from src.utils.path import PathUtil

# 获取应用配置
settings = get_app_settings()

# 创建orm同步引擎
sync_engine = create_engine(
    settings.master_sync_orm_url,
    echo=settings.orm_echo,
    future=settings.orm_future,
    pool_size=settings.orm_process_pool_size,
    pool_pre_ping=settings.orm_pool_pre_ping,
)

# 获取alembic配置
alembic_config = Config(PathUtil.get_alembic_ini_path().as_posix())

# 模型基类必须引入所有相关的模型到同一个文件才能进行创建
# 通过locals进行变量动态添加
locals_var = locals()
# 获取数据模型的名称与类对象的生成器
dynamic_models = AutoImportUtil.auto_import_orm_model(
    PathUtil.get_project_root_path(), PathUtil.get_business_path()
)

for model_name, model_class in dynamic_models:
    # 动态添加变量
    locals_var[model_name] = model_class


class OrmOperate:
    """cli的ORM相关操作类"""

    @staticmethod
    def alembic_init_orm(msg: str = "init"):
        """通过alembic初始化数据库"""
        # 确保alembic的versions文件夹已正常创建
        OrmOperate.alembic_mkdir_versions()
        # 使用alembic的revision指令构建迁移脚本
        command.revision(alembic_config, autogenerate=True, message=msg)
        # 使用alembic的upgrade指令进行迁移操作
        command.upgrade(alembic_config, revision="head")

    @staticmethod
    def alembic_drop_orm():
        """通过alembic让数据库回到初始状态"""
        # 清空alembic的versions相关信息
        OrmOperate.alembic_clear_versions()

    @staticmethod
    def alembic_revision(msg: str):
        """alembic生成迁移文件"""
        command.revision(alembic_config, autogenerate=True, message=msg)

    @staticmethod
    def alembic_upgrade(revision: str):
        """alembic迁移"""
        command.upgrade(alembic_config, revision=revision)

    @staticmethod
    def alembic_downgrade(revision: str):
        """alembic版本回退"""
        command.downgrade(alembic_config, revision=revision)

    @staticmethod
    def alembic_mkdir_versions():
        """创建alembic的versions, 已存在则不做任何操作"""
        versions_dir_path = PathUtil.get_alembic_versions_dir_path()

        if not versions_dir_path.exists():
            versions_dir_path.mkdir()

    @staticmethod
    def alembic_clear_versions():
        """alembic清空versions相关信息"""
        # 先进行alembic版本回退, 回退到base
        OrmOperate.alembic_downgrade("base")
        # 再清空alembic的versions文件夹(删除后重新创建)
        versions_dir_path = PathUtil.get_alembic_versions_dir_path()
        shutil.rmtree(versions_dir_path)
        versions_dir_path.mkdir()
